import time
import data_utils
import numpy as np
import pandas as pd
import os
import matplotlib.pyplot as plt
from dtw import dtw, rabinerJuangStepPattern
from tqdm import tqdm


def DTW_data_analysis():
    # read from xlsx
    data = pd.read_excel("output/DTW/distance_matrix.xlsx", index_col=0)
    # 获取所有C-C的10x10的矩阵
    # 在name=A,B,C的三种情况，分别绘制子图

    fig, axs = plt.subplots(3, figsize=(7, 15))
    names = ["A", "B", "C"]
    for idx, name in enumerate(names):
        sub_data = data.loc[f"{name}1":f"{name}10", "C1":"C10"]
        ground_truth_matrix = data_utils.read_ground_truth(1, type="matrix")
        target_values = []
        for i in range(10):
            for j in range(10):
                if i == j:
                    continue
                if ground_truth_matrix[i][j] == 1:
                    target_values.append(sub_data[f"C{j+1}"][f"{name}{i+1}"])
        sub_data = sub_data.values.flatten()
        sub_data = sorted(sub_data)
        axs[idx].plot(sub_data)
        for tar in target_values:
            axs[idx].scatter(sub_data.index(tar), tar, c="r")
        axs[idx].title.set_text(f"{name} -> C")
    plt.savefig("output/tmp/distance_matrix.png")


def dtw_process():
    start_time = time.time()
    data = data_utils.read_Dream4_time_series_data(1)
    data = pd.DataFrame(data, columns=data_utils.DataName.D3ColumnNames)
    distances_matrix = [[0] * 30 for _ in range(30)]
    for i, name_1 in enumerate(tqdm(data.columns)):
        for j, name_2 in enumerate(data.columns):
            if name_1 == name_2:
                continue
            dir_name = f"output/DTW/{name_1}"
            if not os.path.exists(dir_name):
                os.makedirs(dir_name)
            alignment = dtw(data[name_1], data[name_2], keep_internals=True)
            cur_distance = alignment.normalizedDistance
            distances_matrix[i][j] = cur_distance
            alignment.plot(type="threeway")
            plt.savefig(f"output/DTW/{name_1}/{name_2}_threeway.png")
            plt.close()
            dtw(
                data[name_1],
                data[name_2],
                keep_internals=True,
                step_pattern=rabinerJuangStepPattern(6, "c"),
            ).plot(type="twoway", offset=-2)
            plt.savefig(f"output/DTW/{name_1}/{name_2}_twoway.png")
            plt.close()
            # break
        # break
    print(distances_matrix)
    # 将distance_matrix这个矩阵转为dataframe，保存到xlsx文件中
    distances_matrix = pd.DataFrame(
        distances_matrix, columns=data.columns, index=data.columns
    )
    distances_matrix.to_excel("output/DTW/distance_matrix.xlsx")
    print(f"cost time: {time.time() - start_time}")


if __name__ == "__main__":
    dtw_process()
